package com.fmm.medium;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class HighFrequencyElement {

    /**
     * @param nums 出现的数字集合
     * @param k 出现频率最高的 k 个元素
     * @return 出现频率最高的 k 个元素
     */
    public static int[] topKFrequent(int[] nums, int k) {
        // 记录数字的出现频率
        Map<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            map.put(nums[i], map.getOrDefault(nums[i], 0) + 1);
        }
        // 小顶堆，记录频率最高的 k 个元素
        PriorityQueue<Frequency> pq = new PriorityQueue<>(k, Comparator.comparingInt(Frequency::freq));
        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            // 如果小顶堆的大小比 k 小，直接加入，如果小顶堆的大小超过 k，比较堆顶元素和当前元素，更小的加入
            if (pq.size() < k) {
                pq.offer(new Frequency(entry.getKey(), entry.getValue()));
            } else {
                Frequency smallest = pq.peek();
                if (smallest.freq() < entry.getValue()) {
                    pq.poll();
                    pq.offer(new Frequency(entry.getKey(), entry.getValue()));
                }
            }
        }
        int[] res = new int[k];
        // 获取小顶堆内所有的元素
        for (int i = 0; i < k; i++) {
            res[i] = pq.poll().num();
        }
        return res;
    }

    private record Frequency(int num, int freq) {
    }

    public static void main(String[] args) {
        int[] nums = {4, 1, -1, 2, -1, 2, 3};
        int k = 2;
        int[] res = topKFrequent(nums, k);
        for (int i : res) {
            System.out.print(i + " ");
        }
    }
}
